from jinja2 import Template
import pycuda.autoinit
import pycuda.driver as drv
from pycuda.compiler import SourceModule

# tpl = Template("""
#     __global__ void add({{ type_name }} *tgt,{{ type_name }} *op1,{{ type_name }} *op2)
#     {
#       int idx = threadIdx.x + {{ thread_block_size }} * {{block_size}} * blockIdx.x;
#       {% for i in range(block_size) %}
#           {% set offset = i*thread_block_size %}
#           tgt[idx + {{ offset }}] = op1[idx + {{ offset }}] + op2[idx + {{ offset }}];
#       {% endfor %}
#     }""")
fo = open("nn_gpu.cu", "rb")
content = str(fo.read(), encoding = "utf8")
tpl = Template(content)

rendered_tpl = tpl.render(block_size=32)

print(rendered_tpl)

mod = SourceModule(rendered_tpl)